# edited by levi    @滚球系统视觉
# 卡尔曼滤波，速度模型预测，防止色块跳动丢失目标

from maix import image, camera, display, uart, gpio, pinmap,app
import struct
import numpy as np

pinmap.set_pin_function("B3", "GPIOB3")
led = gpio.GPIO("GPIOB3", gpio.Mode.OUT)
led.value(1)

# 带有速度模型的卡尔曼滤波器类
class KalmanFilterWithVelocity:
    def __init__(self, process_noise=1e-1, sensor_noise=1e-1, estimated_error=0.0):
        self.process_noise = process_noise
        self.sensor_noise = sensor_noise
        self.estimated_error = estimated_error
        self.state = np.array([0.0, 0.0])  # 初始状态 [位置, 速度]
        self.covariance = np.eye(2)  # 初始协方差矩阵

    def predict(self):
        delta_t = 1  # 时间步长
        F = np.array([[1, delta_t], [0, 1]])  # 状态转移矩阵
        self.state = F.dot(self.state)  # 状态预测
        Q = self.process_noise * np.eye(2)  # 过程噪声协方差矩阵
        self.covariance = F.dot(self.covariance).dot(F.T) + Q  # 协方差预测

    def update(self, measurement):
        if measurement is not None:
            H = np.array([[1, 0]])  # 观测矩阵
            R = self.sensor_noise * np.eye(1)  # 观测噪声协方差矩阵
            K = self.covariance.dot(H.T).dot(np.linalg.inv(H.dot(self.covariance).dot(H.T) + R))  # 卡尔曼增益
            self.state = self.state + K.dot(measurement - H.dot(self.state))  # 状态更新
            self.covariance = (np.eye(2) - K.dot(H)).dot(self.covariance)  # 协方差更新

    def get_state(self):
        return self.state[0]

    def get_velocity(self):
        return self.state[1]

# 初始化卡尔曼滤波器
kf_x = KalmanFilterWithVelocity()
kf_y = KalmanFilterWithVelocity()

device = "/dev/ttyS0"
serial0 = uart.UART(device, 115200)
cam = camera.Camera(640, 480)
disp = display.Display()
thresholds = [[67, 81, -9, 9, -12, 8]]
cx, cy, x, y, w, h = 0, 0, 0, 0, 0, 0
ccx, ccy = 170 + (360 / 2), 40 + (350 / 2)

def send_four_uchar(c1, c2):
    sentdata = struct.pack(">BhhB",  # 大端模式
                           0xAB,
                           int(c1),  
                           int(c2), 
                           0xBA
                           )
    serial0.write(sentdata)
    print(sentdata)
    unpacked_data = struct.unpack(">BhhB", sentdata)
    print(unpacked_data)

while not app.need_exit():
    img = cam.read()
    img.draw_rect(170, 40, 360, 350, image.Color.from_rgb(0, 255, 0), 2)
    objsholds = [thresholds[0]]
    blobs = img.find_blobs(objsholds, area_threshold=200, pixels_threshold=300, merge=True, roi=[170, 40, 360, 350])
    if blobs:
        for blob in blobs:
            img.draw_rect(blob[0], blob[1], blob[2], blob[3], image.COLOR_RED)
            cx = blob[0] + blob[2] / 2
            cy = blob[1] + blob[3] / 2
            x = blob[0]
            y = blob[1]
            w = blob[2]
            h = blob[3]
    else:
        cx, cy = None, None  # 如果没有找到色块，设置测量值为None
    
    # 预测步骤
    kf_x.predict()
    kf_y.predict()
    
    # 更新步骤
    kf_x.update(cx)
    kf_y.update(cy)
    
    # 获取滤波后的状态
    filtered_cx = kf_x.get_state()
    filtered_cy = kf_y.get_state()
    
    # 获取滤波后的速度
    velocity_x = kf_x.get_velocity()
    velocity_y = kf_y.get_velocity()
    
    img.draw_cross(int(filtered_cx), int(filtered_cy), image.Color.from_rgb(255, 0, 0), size=10, thickness=2)
    img.draw_string(10, 10, "{:.2f} {:.2f}".format(filtered_cx, filtered_cy), image.Color.from_rgb(255, 0, 0), scale=3)
    img.draw_string(10, 400, "Vx: {:.2f} Vy: {:.2f}".format(velocity_x, velocity_y), image.Color.from_rgb(0, 255, 0), scale=3) #速度

    if filtered_cx < ccx:
        filtered_cx = (ccx - filtered_cx)
    elif filtered_cx > ccx:
        filtered_cx = -(filtered_cx - ccx)
    else:
        filtered_cx = 0

    if filtered_cy < ccy:
        filtered_cy = ccy - filtered_cy
    elif filtered_cy > ccy:
        filtered_cy = -(filtered_cy - ccy)
    else:
        filtered_cy = 0
    
    if filtered_cx>=180:
        filtered_cx=180
    elif filtered_cx<=-180:
        filtered_cx=-180
    if filtered_cy>=175:
        filtered_cy=175
    elif filtered_cy<=-175:
        filtered_cy=-175
        
    # img.draw_string(10, 400, "{:.2f} {:.2f}".format(filtered_cx, filtered_cy), image.Color.from_rgb(0, 255, 0), scale=3)
    img.draw_string(10, 440, "{}    {}".format(int(filtered_cx), int(filtered_cy)), image.Color.from_rgb(0, 255, 0), scale=3)
    send_four_uchar(filtered_cx, filtered_cy)

    disp.show(img)
